# %%
from conn import ac_task_col
from tqdm import tqdm

# %%
anno_type = 31

total = ac_task_col.count_documents({"annotate_type": anno_type})


# %%

print(total)

# %%
ids = []

for task in tqdm(
    ac_task_col.find(
        {"annotate_type": anno_type}, {"clip_ids": 1, "clip_annotated_records": 1}
    ),
    total=total,
):
    if len(task["clip_ids"]) != len(task["clip_annotated_records"]):
        ids.append(task["_id"])
        print(task["_id"])

# %%
print(len(ids))

# %%
for task in ac_task_col.find(
    {"_id": {"$in": ids}}, {"clip_ids": 1, "clip_annotated_records": 1}
):
    clip_record_map = {r["clip_id"]: r for r in task["clip_annotated_records"]}
    new_records = []
    for clip_id in task["clip_ids"]:
        new_records.append(clip_record_map[clip_id])

    assert len(new_records) == len(task["clip_ids"])
    ac_task_col.update_one(
        {"_id": task["_id"]},
        {"$set": {"clip_annotated_records": new_records}},
    )
